#include "Rijndael.h"
#include "lang/System.h"
#include "array/Arrays.h"
#include "GALLOIS.h"
using namespace yzrilyzr_lang;
using namespace yzrilyzr_array;
namespace yzrilyzr_crypto{
	KeyExpander::KeyExpander(const ByteArray & key){
		nk=key.length / 4; // 4, 6 or 8 for 128, 192 or 256-bit keys
		nr=nk + 6; // 10, 12 or 14 rounds for 128, 192 or 256-bit keys
		keyExpansion(key);
	}
	void KeyExpander::keyExpansion(const ByteArray & initialKey){
		int32_t n=4 * (nr + 1); // total number of 32-bit words in expanded key
		ByteArray schedule(4 * n);
		// Copy the initial key
		System::arraycopy(initialKey, 0, schedule, 0, initialKey.length);
		int32_t bytesGenerated=initialKey.length;
		int32_t rconIteration=1;
		ByteArray temp(4);
		while(bytesGenerated < 4 * n){
			System::arraycopy(schedule, bytesGenerated - 4, temp, 0, 4);
			if(bytesGenerated % (4 * nk) == 0){
				keyExpansionCore(temp, rconIteration);
				rconIteration++;
			}
			// For 256-bit key, apply s-box to all bytes every 32 bytes
			if(nk > 6 && bytesGenerated % (4 * nk) == 16){
				for(int32_t i=0;i < 4;i++){
					temp[i]=(int8_t)GALLOIS::sBox[toUnsignedInt(temp[i])];
				}
			}
			// XOR temp with the word nk positions back
			for(int32_t i=0;i < 4;i++){
				schedule[bytesGenerated + i]=(int8_t)(temp[i] ^ schedule[bytesGenerated + i - 4 * nk]);
			}
			bytesGenerated+=4;
		}
		// Split into round keys (each 16 bytes)
		this->keys=Array<ByteArray>(nr + 1);
		for(int32_t i=0;i <= nr;i++){
			this->keys[i]=ByteArray(16);
			System::arraycopy(schedule, 16 * i, this->keys[i], 0, 16);
		}
	}
	void KeyExpander::keyExpansionCore(const ByteArray & in, int32_t round){
		// Rotate word
		int8_t tmp=in[0];
		in[0]=in[1];
		in[1]=in[2];
		in[2]=in[3];
		in[3]=tmp;
		// Sub word
		for(int32_t i=0;i < 4;i++){
			in[i]=(int8_t)GALLOIS::sBox[toUnsignedInt(in[i])];
		}
		// Rcon
		in[0]^=(int8_t)GALLOIS::rcon[round];
	}
	State::State() :data(16){}
	void State::setKey(const KeyExpander & ke){
		this->rounds=ke.keys.length - 1;
		this->keys=ke.keys;
	}

	void State::input(const ByteArray & data){
		this->round=0;
		apply(data);
	}
	void State::apply(const ByteArray & data){
		System::arraycopy(data, 0, this->data, 0, 16);
	}
	const ByteArray State::getKey(int32_t i){
		if(i >= 0 && i < keys.length) return keys[i];
		else throw IndexOutOfBoundsException("key out of bounds");
	}
	ByteArray State::getData(){
		return this->data;
	}
	void State::addRoundKey(const ByteArray & key){
		for(int32_t i=0;i < 16;i++){
			this->data[i]^=key[i];
		}
	}
	void State::subBytes(){
		for(int32_t i=0;i < 16;i++){
			this->data[i]=(int8_t)GALLOIS::sBox[toUnsignedInt(this->data[i])];
		}
	}
	void State::InverseSubBytes(){
		for(int32_t i=0;i < 16;i++){
			this->data[i]=(int8_t)GALLOIS::invsBox[toUnsignedInt(this->data[i])];
		}
	}
	void State::shiftRows(){
		ByteArray shift(16);
		static constexpr int8_t index[]={
			0, 5, 10, 15, 4, 9, 14, 3, 8, 13, 2, 7, 12, 1, 6, 11
		};
		for(int32_t i=0;i < 16;i++){
			shift[i]=this->data[index[i]];
		}
		this->apply(shift);
	}
	void State::InverseShiftRows(){
		ByteArray shift(16);
		static constexpr int8_t index[]={
			0, 13, 10, 7, 4, 1, 14, 11, 8, 5, 2, 15, 12, 9, 6, 3
		};
		for(int32_t i=0;i < 16;i++){
			shift[i]=this->data[index[i]];
		}
		this->apply(shift);
	}
	void State::mixColumns(){
		ByteArray mix(16);
		for(int32_t col=0;col < 4;col++){
			int32_t idx0=4 * col;
			int32_t idx1=1 + 4 * col;
			int32_t idx2=2 + 4 * col;
			int32_t idx3=3 + 4 * col;
			mix[idx0]=(int8_t)(GALLOIS::mul2[toUnsignedInt(this->data[idx0])] ^
							   GALLOIS::mul3[toUnsignedInt(this->data[idx1])] ^
							   this->data[idx2] ^
							   this->data[idx3]);
			mix[idx1]=(int8_t)(this->data[idx0] ^
							   GALLOIS::mul2[toUnsignedInt(this->data[idx1])] ^
							   GALLOIS::mul3[toUnsignedInt(this->data[idx2])] ^
							   this->data[idx3]);
			mix[idx2]=(int8_t)(this->data[idx0] ^
							   this->data[idx1] ^
							   GALLOIS::mul2[toUnsignedInt(this->data[idx2])] ^
							   GALLOIS::mul3[toUnsignedInt(this->data[idx3])]);
			mix[idx3]=(int8_t)(GALLOIS::mul3[toUnsignedInt(this->data[idx0])] ^
							   this->data[idx1] ^
							   this->data[idx2] ^
							   GALLOIS::mul2[toUnsignedInt(this->data[idx3])]);
		}
		this->apply(mix);
	}
	void State::InverseMixColumns(){
		ByteArray mix(16);
		for(int32_t col=0;col < 4;col++){
			int32_t idx0=4 * col;
			int32_t idx1=1 + 4 * col;
			int32_t idx2=2 + 4 * col;
			int32_t idx3=3 + 4 * col;
			mix[idx0]=(int8_t)(GALLOIS::mul14[toUnsignedInt(this->data[idx0])] ^
							   GALLOIS::mul11[toUnsignedInt(this->data[idx1])] ^
							   GALLOIS::mul13[toUnsignedInt(this->data[idx2])] ^
							   GALLOIS::mul9[toUnsignedInt(this->data[idx3])]);
			mix[idx1]=(int8_t)(GALLOIS::mul9[toUnsignedInt(this->data[idx0])] ^
							   GALLOIS::mul14[toUnsignedInt(this->data[idx1])] ^
							   GALLOIS::mul11[toUnsignedInt(this->data[idx2])] ^
							   GALLOIS::mul13[toUnsignedInt(this->data[idx3])]);
			mix[idx2]=(int8_t)(GALLOIS::mul13[toUnsignedInt(this->data[idx0])] ^
							   GALLOIS::mul9[toUnsignedInt(this->data[idx1])] ^
							   GALLOIS::mul14[toUnsignedInt(this->data[idx2])] ^
							   GALLOIS::mul11[toUnsignedInt(this->data[idx3])]);
			mix[idx3]=(int8_t)(GALLOIS::mul11[toUnsignedInt(this->data[idx0])] ^
							   GALLOIS::mul13[toUnsignedInt(this->data[idx1])] ^
							   GALLOIS::mul9[toUnsignedInt(this->data[idx2])] ^
							   GALLOIS::mul14[toUnsignedInt(this->data[idx3])]);
		}
		this->apply(mix);
	}
	yzrilyzr_array::ByteArray Rijndael::encrypt(const yzrilyzr_array::ByteArray & in){
		state.input(in);
		encryptCore();
		return state.getData();
	}
	yzrilyzr_array::ByteArray Rijndael::decrypt(const yzrilyzr_array::ByteArray & in){
		state.input(in);
		decryptCore();
		return state.getData();
	}
	void Rijndael::decryptCore(){
		// Reverse of encryptCore
		int32_t numberOfRounds=state.rounds;
		state.addRoundKey(state.getKey(state.rounds));
		for(state.round=numberOfRounds - 1;state.round > 0;--state.round){
			state.InverseShiftRows();
			state.InverseSubBytes();
			state.addRoundKey(state.getKey(state.round));
			state.InverseMixColumns();
		}
		// Final round
		state.InverseShiftRows();
		state.InverseSubBytes();
		state.addRoundKey(state.getKey(state.round));
	}
	void Rijndael::encryptCore(){
		// Set to round 0
		state.round=0;
		int32_t numberOfRounds=state.rounds;
		// Initial round, key expansion already happened
		state.addRoundKey(state.getKey(state.round));
		// The remaining rounds		
		for(state.round=1;(state.round) < numberOfRounds; ++(state.round)){
 			state.subBytes();
			state.shiftRows();
			state.mixColumns();
			state.addRoundKey(state.getKey(state.round));
		}
		// Final round
		state.subBytes();
		state.shiftRows();
		state.addRoundKey(state.getKey(state.round));
	}
	void Rijndael::setKey(const FAQKey & rawKey){
		setKey(rawKey.getEncoded());
	}
	void Rijndael::setKey(const yzrilyzr_array::ByteArray & rawKey){
		int32_t keyLength=rawKey.length;
		if(keyLength != 16 && keyLength != 24 && keyLength != 32){
			throw IllegalArgumentException(
				"Rijndael key must be 16, 24, or 32 bytes, current is " +
				keyLength);
		}
		KeyExpander ke(rawKey);
		state.setKey(ke);
	}
}